import pandas as pd
import matplotlib.pyplot as plt
import mplfinance as mpf
import sympy as sp
import os

# 读取CSV文件
file_path = os.path.join('csv','123227_full.csv')
csv_data = pd.read_csv(file_path)

# 将日期列转换为datetime类型
csv_data['date'] = pd.to_datetime(csv_data['date'])

# 设置日期列为索引
csv_data.set_index('date', inplace=True)

# 绘制K线图和成交量柱状图，设置上涨和下跌颜色
mc = mpf.make_marketcolors(up='r', down='g')  # 设置上涨颜色为红色，下跌颜色为绿色
s = mpf.make_mpf_style(marketcolors=mc)

kwargs = dict(
    type='candle', 
    volume=True, 
    figratio=(12, 8), 
    figscale=1.5,
    datetime_format='%Y-%m-%d',
    style=s  # 使用自定义样式
)

# 表达式解析函数
def parse_expression(expression, data):
    # 手动创建所有需要的符号
    symbols = {col: sp.symbols(col) for col in data.columns}
    expr = sp.sympify(expression, locals=symbols)
    print(f"expr type: {type(expr)}, expr value: {expr}")
    return data.apply(lambda row: float(expr.subs({symbols[col]: row[col] for col in data.columns})), axis=1)

# 时间筛选函数
def filter_data_by_time(data, start_date, end_date):
    return data.loc[(data.index >= start_date) & (data.index <= end_date)]

# 用户输入表达式和时间范围
user_expression = input('请输入线形图对应的值的表达式（例如：open + close）：')
start_date_str = "2024-01-10"#input('请输入开始日期（YYYY-MM-DD）：')
end_date_str = "2025-03-10"#input('请输入结束日期（YYYY-MM-DD）：')

start_date = pd.to_datetime(start_date_str)
end_date = pd.to_datetime(end_date_str)

# 筛选数据
filtered_data = filter_data_by_time(csv_data, start_date, end_date)

# 解析表达式
line_data = parse_expression(user_expression, filtered_data)

ap = mpf.make_addplot(line_data, color='r')

# 绘制图形
fig, axes = mpf.plot(filtered_data, **kwargs, addplot=ap, xrotation=45)

# 获取主图的坐标轴
ax = axes[0]
# 设置日期格式
import matplotlib.dates as mdates
ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))

plt.show()